"""
Alpha101基础运算符
================

实现Alpha101因子计算所需的所有基础运算符，包括：
- 时间序列运算符（ts_rank, ts_sum, ts_mean等）
- 截面运算符（rank, scale, neutralize等）
- 相关性运算符（correlation, covariance等）

所有运算符都使用Numba JIT编译进行加速优化。
"""

import numpy as np
import numba
from numba import jit, prange
from scipy.stats import rankdata
import warnings

# 忽略numba的性能警告
warnings.filterwarnings('ignore', category=numba.NumbaPerformanceWarning)


class BaseOperators:
    """基础运算符类，包含所有Alpha101所需的基础计算函数"""
    
    @staticmethod
    @jit(nopython=True, parallel=True)
    def ts_sum(data, window):
        """
        时间序列求和
        
        Parameters:
        -----------
        data : np.ndarray, shape (n_stocks, n_days)
            输入数据
        window : int
            时间窗口长度
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            滚动求和结果
        """
        n_stocks, n_days = data.shape
        result = np.empty((n_stocks, n_days))
        result[:, :window-1] = np.nan
        
        for i in prange(n_stocks):
            for j in range(window-1, n_days):
                result[i, j] = np.sum(data[i, j-window+1:j+1])
        
        return result
    
    @staticmethod
    @jit(nopython=True, parallel=True)
    def ts_mean(data, window):
        """
        时间序列均值
        
        Parameters:
        -----------
        data : np.ndarray, shape (n_stocks, n_days)
            输入数据
        window : int
            时间窗口长度
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            滚动均值结果
        """
        n_stocks, n_days = data.shape
        result = np.empty((n_stocks, n_days))
        result[:, :window-1] = np.nan
        
        for i in prange(n_stocks):
            for j in range(window-1, n_days):
                result[i, j] = np.mean(data[i, j-window+1:j+1])
        
        return result
    
    @staticmethod
    @jit(nopython=True, parallel=True)
    def ts_std(data, window):
        """
        时间序列标准差
        
        Parameters:
        -----------
        data : np.ndarray, shape (n_stocks, n_days)
            输入数据
        window : int
            时间窗口长度
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            滚动标准差结果
        """
        n_stocks, n_days = data.shape
        result = np.empty((n_stocks, n_days))
        result[:, :window-1] = np.nan
        
        for i in prange(n_stocks):
            for j in range(window-1, n_days):
                window_data = data[i, j-window+1:j+1]
                result[i, j] = np.std(window_data)
        
        return result
    
    @staticmethod
    @jit(nopython=True, parallel=True)
    def ts_rank(data, window):
        """
        时间序列排序（当前值在窗口内的排名）
        
        Parameters:
        -----------
        data : np.ndarray, shape (n_stocks, n_days)
            输入数据
        window : int
            时间窗口长度
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            排名结果（0-1之间）
        """
        n_stocks, n_days = data.shape
        result = np.empty((n_stocks, n_days))
        result[:, :window-1] = np.nan
        
        for i in prange(n_stocks):
            for j in range(window-1, n_days):
                window_data = data[i, j-window+1:j+1]
                current_value = data[i, j]
                rank = np.sum(window_data <= current_value) / window
                result[i, j] = rank
        
        return result
    
    @staticmethod
    @jit(nopython=True, parallel=True)
    def ts_argmax(data, window):
        """
        时间序列最大值位置
        
        Parameters:
        -----------
        data : np.ndarray, shape (n_stocks, n_days)
            输入数据
        window : int
            时间窗口长度
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            最大值位置（0到window-1）
        """
        n_stocks, n_days = data.shape
        result = np.empty((n_stocks, n_days))
        result[:, :window-1] = np.nan
        
        for i in prange(n_stocks):
            for j in range(window-1, n_days):
                window_data = data[i, j-window+1:j+1]
                result[i, j] = np.argmax(window_data)
        
        return result
    
    @staticmethod
    @jit(nopython=True, parallel=True)
    def ts_argmin(data, window):
        """
        时间序列最小值位置
        
        Parameters:
        -----------
        data : np.ndarray, shape (n_stocks, n_days)
            输入数据
        window : int
            时间窗口长度
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            最小值位置（0到window-1）
        """
        n_stocks, n_days = data.shape
        result = np.empty((n_stocks, n_days))
        result[:, :window-1] = np.nan
        
        for i in prange(n_stocks):
            for j in range(window-1, n_days):
                window_data = data[i, j-window+1:j+1]
                result[i, j] = np.argmin(window_data)
        
        return result
    
    @staticmethod
    @jit(nopython=True, parallel=True)
    def ts_corr(x, y, window):
        """
        时间序列相关性
        
        Parameters:
        -----------
        x, y : np.ndarray, shape (n_stocks, n_days)
            输入数据
        window : int
            时间窗口长度
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            相关系数
        """
        n_stocks, n_days = x.shape
        result = np.empty((n_stocks, n_days))
        result[:, :window-1] = np.nan
        
        for i in prange(n_stocks):
            for j in range(window-1, n_days):
                x_window = x[i, j-window+1:j+1]
                y_window = y[i, j-window+1:j+1]
                
                # 计算相关系数
                x_mean = np.mean(x_window)
                y_mean = np.mean(y_window)
                
                numerator = np.sum((x_window - x_mean) * (y_window - y_mean))
                x_std = np.sqrt(np.sum((x_window - x_mean) ** 2))
                y_std = np.sqrt(np.sum((y_window - y_mean) ** 2))
                
                if x_std > 1e-8 and y_std > 1e-8:
                    result[i, j] = numerator / (x_std * y_std)
                else:
                    result[i, j] = np.nan
        
        return result
    
    @staticmethod
    @jit(nopython=True, parallel=True)
    def delta(data, period=1):
        """
        差分运算
        
        Parameters:
        -----------
        data : np.ndarray, shape (n_stocks, n_days)
            输入数据
        period : int
            差分周期
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            差分结果
        """
        n_stocks, n_days = data.shape
        result = np.empty((n_stocks, n_days))
        result[:, :period] = np.nan
        
        for i in prange(n_stocks):
            for j in range(period, n_days):
                result[i, j] = data[i, j] - data[i, j-period]
        
        return result
    
    @staticmethod
    @jit(nopython=True, parallel=True)
    def delay(data, period=1):
        """
        延迟运算
        
        Parameters:
        -----------
        data : np.ndarray, shape (n_stocks, n_days)
            输入数据
        period : int
            延迟周期
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            延迟结果
        """
        n_stocks, n_days = data.shape
        result = np.empty((n_stocks, n_days))
        result[:, :period] = np.nan
        
        for i in prange(n_stocks):
            for j in range(period, n_days):
                result[i, j] = data[i, j-period]
        
        return result
    
    @staticmethod
    @jit(nopython=True, parallel=True)
    def ts_min(data, window):
        """
        时间序列最小值
        
        Parameters:
        -----------
        data : np.ndarray, shape (n_stocks, n_days)
            输入数据
        window : int
            时间窗口长度
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            滚动最小值结果
        """
        n_stocks, n_days = data.shape
        result = np.empty((n_stocks, n_days))
        result[:, :window-1] = np.nan
        
        for i in prange(n_stocks):
            for j in range(window-1, n_days):
                window_data = data[i, j-window+1:j+1]
                result[i, j] = np.min(window_data)
        
        return result
    
    @staticmethod
    @jit(nopython=True, parallel=True)
    def ts_max(data, window):
        """
        时间序列最大值
        
        Parameters:
        -----------
        data : np.ndarray, shape (n_stocks, n_days)
            输入数据
        window : int
            时间窗口长度
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            滚动最大值结果
        """
        n_stocks, n_days = data.shape
        result = np.empty((n_stocks, n_days))
        result[:, :window-1] = np.nan
        
        for i in prange(n_stocks):
            for j in range(window-1, n_days):
                window_data = data[i, j-window+1:j+1]
                result[i, j] = np.max(window_data)
        
        return result
    
    @staticmethod
    @jit(nopython=True, parallel=True)
    def ts_cov(x, y, window):
        """
        时间序列协方差
        
        Parameters:
        -----------
        x, y : np.ndarray, shape (n_stocks, n_days)
            输入数据
        window : int
            时间窗口长度
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            协方差结果
        """
        n_stocks, n_days = x.shape
        result = np.empty((n_stocks, n_days))
        result[:, :window-1] = np.nan
        
        for i in prange(n_stocks):
            for j in range(window-1, n_days):
                x_window = x[i, j-window+1:j+1]
                y_window = y[i, j-window+1:j+1]
                
                # 计算协方差
                x_mean = np.mean(x_window)
                y_mean = np.mean(y_window)
                
                covariance = np.sum((x_window - x_mean) * (y_window - y_mean)) / (window - 1)
                result[i, j] = covariance
        
        return result


class CrossSectionOps:
    """截面运算符类，处理股票间的横截面计算"""
    
    @staticmethod
    def rank(data):
        """
        截面排序
        
        Parameters:
        -----------
        data : np.ndarray, shape (n_stocks, n_days)
            输入数据
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            排名结果（0-1之间）
        """
        n_stocks, n_days = data.shape
        result = np.empty((n_stocks, n_days))
        
        for j in range(n_days):
            day_data = data[:, j]
            valid_mask = ~np.isnan(day_data)
            
            if np.sum(valid_mask) > 1:
                # 使用scipy的rankdata进行排序
                ranks = rankdata(day_data[valid_mask], method='average')
                # 标准化到0-1之间
                normalized_ranks = (ranks - 1) / (len(ranks) - 1)
                result[valid_mask, j] = normalized_ranks
                result[~valid_mask, j] = np.nan
            else:
                result[:, j] = np.nan
        
        return result
    
    @staticmethod
    def scale(data, scale_factor=1.0):
        """
        截面标准化
        
        Parameters:
        -----------
        data : np.ndarray, shape (n_stocks, n_days)
            输入数据
        scale_factor : float
            缩放因子
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            标准化结果
        """
        n_stocks, n_days = data.shape
        result = np.empty((n_stocks, n_days))
        
        for j in range(n_days):
            day_data = data[:, j]
            valid_mask = ~np.isnan(day_data)
            
            if np.sum(valid_mask) > 1:
                valid_data = day_data[valid_mask]
                mean_val = np.mean(valid_data)
                std_val = np.std(valid_data)
                
                if std_val > 1e-8:
                    result[valid_mask, j] = (valid_data - mean_val) / std_val * scale_factor
                else:
                    result[valid_mask, j] = 0.0
                result[~valid_mask, j] = np.nan
            else:
                result[:, j] = np.nan
        
        return result
    
    @staticmethod
    def neutralize(data, industry_codes):
        """
        行业中性化
        
        Parameters:
        -----------
        data : np.ndarray, shape (n_stocks, n_days)
            输入数据
        industry_codes : np.ndarray, shape (n_stocks,)
            行业代码
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            中性化结果
        """
        n_stocks, n_days = data.shape
        result = np.empty((n_stocks, n_days))
        
        unique_industries = np.unique(industry_codes[~np.isnan(industry_codes)])
        
        for j in range(n_days):
            day_data = data[:, j]
            result[:, j] = day_data.copy()
            
            for industry in unique_industries:
                industry_mask = (industry_codes == industry) & (~np.isnan(day_data))
                
                if np.sum(industry_mask) > 1:
                    industry_data = day_data[industry_mask]
                    industry_mean = np.mean(industry_data)
                    result[industry_mask, j] = industry_data - industry_mean
        
        return result
    
    @staticmethod
    def winsorize(data, limits=(0.025, 0.975)):
        """
        截面去极值
        
        Parameters:
        -----------
        data : np.ndarray, shape (n_stocks, n_days)
            输入数据
        limits : tuple
            上下分位数限制
            
        Returns:
        --------
        np.ndarray, shape (n_stocks, n_days)
            去极值结果
        """
        n_stocks, n_days = data.shape
        result = data.copy()
        
        for j in range(n_days):
            day_data = data[:, j]
            valid_mask = ~np.isnan(day_data)
            
            if np.sum(valid_mask) > 2:
                valid_data = day_data[valid_mask]
                lower_bound = np.percentile(valid_data, limits[0] * 100)
                upper_bound = np.percentile(valid_data, limits[1] * 100)
                
                # 截断极值
                clipped_data = np.clip(valid_data, lower_bound, upper_bound)
                result[valid_mask, j] = clipped_data
        
        return result


# 辅助函数
@jit(nopython=True)
def power(data, exp):
    """幂运算"""
    return np.power(data, exp)


@jit(nopython=True)
def log(data):
    """对数运算"""
    return np.log(np.maximum(data, 1e-8))  # 避免log(0)


@jit(nopython=True)
def abs_func(data):
    """绝对值运算"""
    return np.abs(data)


@jit(nopython=True)
def sign(data):
    """符号函数"""
    return np.sign(data)


@jit(nopython=True)
def max_func(x, y):
    """最大值函数"""
    return np.maximum(x, y)


@jit(nopython=True)
def min_func(x, y):
    """最小值函数"""
    return np.minimum(x, y)


@jit(nopython=True, parallel=True)
def decay_linear(data, window):
    """
    线性衰减权重
    
    Parameters:
    -----------
    data : np.ndarray, shape (n_stocks, n_days)
        输入数据
    window : int
        时间窗口长度
        
    Returns:
    --------
    np.ndarray, shape (n_stocks, n_days)
        线性衰减加权结果
    """
    n_stocks, n_days = data.shape
    result = np.empty((n_stocks, n_days))
    result[:, :window-1] = np.nan
    
    # 创建线性衰减权重：最新的权重最大
    weights = np.arange(1, window + 1, dtype=np.float64)
    weights = weights / np.sum(weights)
    
    for i in prange(n_stocks):
        for j in range(window-1, n_days):
            window_data = data[i, j-window+1:j+1]
            result[i, j] = np.sum(window_data * weights)
    
    return result


@jit(nopython=True)
def product(data, window):
    """
    时间序列乘积
    
    Parameters:
    -----------
    data : np.ndarray, shape (n_stocks, n_days)
        输入数据
    window : int
        时间窗口长度
        
    Returns:
    --------
    np.ndarray, shape (n_stocks, n_days)
        滚动乘积结果
    """
    n_stocks, n_days = data.shape
    result = np.empty((n_stocks, n_days))
    result[:, :window-1] = np.nan
    
    for i in range(n_stocks):
        for j in range(window-1, n_days):
            window_data = data[i, j-window+1:j+1]
            result[i, j] = np.prod(window_data)
    
    return result


def indneutralize(data, industry_codes):
    """
    行业中性化（简化版本，实际应该更复杂）
    
    Parameters:
    -----------
    data : np.ndarray, shape (n_stocks, n_days)
        输入数据
    industry_codes : np.ndarray, shape (n_stocks,)
        行业代码
        
    Returns:
    --------
    np.ndarray, shape (n_stocks, n_days)
        行业中性化结果
    """
    # 这里使用CrossSectionOps的neutralize方法
    cross_ops = CrossSectionOps()
    return cross_ops.neutralize(data, industry_codes)